import torch
from torch import nn
from torch.nn import functional as F
import os

import sys
sys.path.append("../../py")
from utils import save_tensor_as_text, save_tensor_as_c_header, float_tensor2Q88
from mobilenetv2_mini import mobilenetv2
from fuse_bn import fuse_module

SAVED_MODEL = "../../py/mobilenetv2_on_cifar10_mini.pth"

if __name__ == '__main__':
    net = mobilenetv2()
    net.load_state_dict(torch.load(SAVED_MODEL))
    net = fuse_module(net)
    
    img = torch.randn(3, 32, 32)
    out = net(img).reshape([10, 1, 1])
    
    save_tensor_as_text(float_tensor2Q88(out), "img_out_py.txt")
    save_tensor_as_c_header(float_tensor2Q88(img), "img_in")
    save_tensor_as_c_header(float_tensor2Q88(out), "img_out_ref")
        